#!/bin/bash

# 脚本名称：format_disk.sh

# 功能说明：
# 本脚本用于格式化虚拟磁盘或物理磁盘。

# 主要流程包括：
# 1. 安装必要的依赖工具（如 qemu-img、qemu-nbd、parted 等）。
# 2. 根据用户输入判断设备是物理磁盘还是虚拟磁盘镜像。
# 3. 如果是虚拟磁盘镜像，则使用 qemu-nbd 挂载到一个空闲的 nbd 设备上。
# 4. 显示目标设备的详细信息，方便用户核对当前状态。
# 5. 根据用户选择创建分区并格式化为 FAT32 或 EXT4 文件系统。
# 6. 使用 trap 机制在脚本退出时自动断开挂载的 nbd 设备，确保系统资源得到及时释放。

# 注意事项：
# - 需要 root 权限运行部分命令，如 qemu-nbd、parted、mkfs.vfat、mkfs.ext4。
# - 在执行磁盘格式化前，建议备份重要数据，以防万一。

set -euo pipefail

# 安装依赖包
install_dependencies() {
    local required=("qemu-img" "qemu-nbd" "partprobe" "parted" "mkfs.vfat" "mkfs.ext4")
    local missing=()

    for cmd in "${required[@]}"; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            missing+=("$cmd")
        fi
    done

    if [ ${#missing[@]} -eq 0 ]; then
        return 0
    fi

    echo "检测到缺少依赖命令: ${missing[*]}，尝试自动安装..."
    sudo apt-get update -qq
    echo "正在安装 qemu-utils parted e2fsprogs dosfstools..."
    sudo apt-get install -y -qq qemu-utils parted e2fsprogs dosfstools >/dev/null
}

# 全局变量
disk=""
nbd_device=""
target_device=""
virtual_disk_attached=false

# 清理函数：如果虚拟磁盘已挂载，则断开 nbd 设备
cleanup() {
    if $virtual_disk_attached && [ -n "$nbd_device" ]; then
        echo "清理：断开 nbd 设备 $nbd_device"
        sudo qemu-nbd -d "$nbd_device"
    fi
}
trap cleanup EXIT SIGINT SIGTERM

# 查找空闲的 nbd 设备
select_free_nbd() {
    for dev in /dev/nbd*; do
        if [ -z "$(lsblk -n "$dev" | awk 'NR>1')" ]; then
            echo "$dev"
            return 0
        fi
    done
    return 1
}

# 挂载虚拟磁盘
mount_virtual_disk() {
    sudo modprobe nbd max_part=8
    nbd_device=$(select_free_nbd)
    if [ -z "$nbd_device" ]; then
        echo "未找到可用的 nbd 设备，退出。"
        exit 1
    fi
    echo "将虚拟磁盘镜像挂载到 nbd 设备：$nbd_device"
    sudo qemu-nbd -c "$nbd_device" "$disk"
    sudo udevadm settle
    virtual_disk_attached=true
}

# 获取用户输入
get_disk_input() {
    if [ -n "$1" ]; then
        disk="$1"
    else
        echo "没有提供设备或镜像文件，现列出所有磁盘："
        lsblk -o NAME,MAJ:MIN,RM,LABEL,SIZE,FSTYPE,UUID,RO,TYPE,MOUNTPOINT,VENDOR,HOTPLUG -e7
        read -rp "请输入设备路径（如 /dev/sda）或虚拟磁盘镜像文件路径: " disk
    fi
}

# 处理设备类型
process_disk() {
    if [ -b "$disk" ]; then
        echo "选择的物理磁盘: $disk"
        target_device="$disk"
    elif [ -f "$disk" ]; then
        echo "选择的虚拟磁盘镜像文件: $disk"
        mount_virtual_disk
        target_device="$nbd_device"
    else
        echo "错误： '$disk' 既不是块设备也不是普通文件。"
        exit 1
    fi
}

# 显示设备信息
display_disk_info() {
    echo "显示设备信息："
    lsblk -o NAME,MAJ:MIN,RM,LABEL,SIZE,FSTYPE,UUID,RO,TYPE,MOUNTPOINT,VENDOR,HOTPLUG "$target_device"
}

# 更新内核分区表，并稍作等待以确保更新生效
update_partition_table() {
    echo "更新内核分区表..."
    sudo partprobe "$target_device"
    sleep 2
}

# 卸载目标设备上已挂载的所有分区
unmount_partitions() {
    echo "卸载 $target_device 上的所有已挂载分区..."
    local mount_points
    mount_points=$(sudo lsblk -nr -o MOUNTPOINT "$target_device" | grep -v '^$' || true)
    for mp in $mount_points; do
        echo "卸载: $mp"
        sudo umount "$mp" || echo "卸载 $mp 时出现问题"
    done
}

# 根据用户选择创建分区
create_partition_table() {
    local scheme="$1"
    echo "在 $target_device 上初始化 MBR 分区表..."
    sudo parted -s "$target_device" mklabel msdos

    case "$scheme" in
        1)
            echo "创建单个 FAT32 分区..."
            sudo parted -s "$target_device" mkpart primary fat32 0% 100%
            sudo parted -s "$target_device" set 1 boot on
            ;;
        2)
            echo "创建 FAT32 + EXT4 分区..."
            sudo parted -s "$target_device" mkpart primary fat32 0% 1GB
            sudo parted -s "$target_device" set 1 boot on
            sudo parted -s "$target_device" mkpart primary ext4 1GB 100%
            ;;
        3)
            echo "创建 FAT32 + EXT4 + Data 分区..."
            sudo parted -s "$target_device" mkpart primary fat32 0% 1GB
            sudo parted -s "$target_device" set 1 boot on
            sudo parted -s "$target_device" mkpart primary ext4 1GB 50%
            sudo parted -s "$target_device" mkpart primary ext4 50% 100%
            ;;
        *)
            echo "未知的分区方案: $scheme"
            exit 1
            ;;
    esac
}

# 格式化各分区
format_partitions() {
    echo "开始格式化分区..."
    # 使用 lsblk 列出 target_device 下所有类型为 part 的设备
    local partitions
    partitions=$(sudo lsblk -nr -o NAME,TYPE "$target_device" | awk '$2=="part"{print $1}')
    for part in $partitions; do
        # 从分区名中提取分区编号（假设分区名以数字结尾，如 nbd0p1 或 sda1）
        local partnum
        partnum=$(echo "$part" | grep -oE '[0-9]+$')
        case "$partnum" in
            1)
                echo "格式化 /dev/$part 为 FAT32 (BOOT)..."
                sudo mkfs.vfat -F 32 -n BOOT "/dev/$part"
                ;;
            2)
                echo "格式化 /dev/$part 为 EXT4 (rootfs)..."
                echo y | sudo mkfs.ext4 -L rootfs "/dev/$part"
                ;;
            3)
                echo "格式化 /dev/$part 为 EXT4 (userdata)..."
                echo y | sudo mkfs.ext4 -L userdata "/dev/$part"
                ;;
            *)
                echo "发现额外的未知分区: /dev/$part，跳过格式化。"
                ;;
        esac
    done
}

# 主函数
main() {
    install_dependencies
    get_disk_input "$1"
    process_disk
    display_disk_info

    # 用户确认操作
    read -r -p "[WARNING] 此操作将清除 $target_device 上的所有数据，是否继续？ (y/n): " confirm </dev/tty
    if [[ "$confirm" != "y" ]]; then
        echo "操作已取消。"
        exit 0
    fi

    unmount_partitions

    # 选择分区方案
    echo -e "\n请选择分区方案："
    echo " 1) 单个 FAT32 分区"
    echo " 2) FAT32 + EXT4 分区"
    echo " 3) FAT32 + EXT4 + Data 分区"
    read -r -p "请输入方案编号 (1-3): " scheme </dev/tty
    if [[ ! "$scheme" =~ ^[1-3]$ ]]; then
        echo "无效的分区方案选择"
        exit 1
    fi

    create_partition_table "$scheme"
    update_partition_table

    format_partitions
    update_partition_table

    echo "最终 $target_device 的分区布局："
    sudo lsblk -o NAME,FSTYPE,LABEL,SIZE,TYPE "$target_device"

    echo "磁盘准备完成！"
}

main "$@"